package cn.xxi.ai.application.constant;

import cn.xxi.ai.config.OllamaProperties;
import cn.xxi.ai.util.SpringContextHolder;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.googleai.GoogleAiGeminiStreamingChatModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import lombok.Getter;

/**
 * Chat 模型枚举类
 */
@Getter
public enum ChatModel {
    QWEN_2_5_1_5B("qwen2_5:1.5B", "ollama"),
    DEEPSEEK_R1_7B("deepseek-r1:7b", "ollama"),
    QWEN_MAX("qwen-max", "dashscope"),
    GEMINI_2_0_FLASH("gemini-2.0-flash", "gemini"),
    GEMINI_2_0_FLASH_LITE("gemini-2.0-flash-lite", "gemini"),
    GEMINI_1_5_FLASH("gemini-1.5-flash", "gemini"),
    GEMINI_1_5_PRO("gemini-1.5-pro", "gemini");

    private final String modelName;
    private final String provider;

    ChatModel(String modelName, String provider) {
        this.modelName = modelName;
        this.provider = provider;
    }

    public ChatLanguageModel model(Double temperature, Integer topK, Double topP, Double repeatPenalty, Integer seed, Integer numPredict, Integer numCtx) {
        return switch (provider) {
            case "ollama" ->
                    OllamaChatModel.builder().baseUrl(baseUrl()).modelName(modelName).temperature(temperature).topK(topK).topP(topP).repeatPenalty(repeatPenalty).seed(seed).numPredict(numPredict).numCtx(numCtx).logRequests(true).logResponses(true).build();
            case "dashscope" ->
                    QwenChatModel.builder()
                            .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                            .modelName(modelName)
                            .temperature(temperature != null ? temperature.floatValue() : null)
                            .topK(topK).topP(topP)
                            .repetitionPenalty(repeatPenalty != null ? repeatPenalty.floatValue() : null)
                            .seed(seed).build();
            default -> throw new IllegalArgumentException("Unsupported provider: " + provider);
        };
    }

    public StreamingChatLanguageModel stream(Double temperature, Integer topK, Double topP, Double repeatPenalty, Integer seed, Integer numPredict, Integer numCtx) {
        return switch (provider) {
            case "ollama" ->
                    OllamaStreamingChatModel.builder()
                            .baseUrl(baseUrl())
                            .modelName(modelName)
                            .temperature(temperature)
                            .topK(topK)
                            .topP(topP)
                            .repeatPenalty(repeatPenalty)
                            .seed(seed)
                            .numPredict(numPredict)
                            .numCtx(numCtx)
                            .logRequests(true)
                            .logResponses(true)
                            .build();
            case "dashscope" ->
                    QwenStreamingChatModel.builder()
                            .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                            .modelName(modelName)
                            .temperature(temperature != null ? temperature.floatValue() : null)
                            .topK(topK)
                            .topP(topP)
                            .repetitionPenalty(repeatPenalty != null ? repeatPenalty.floatValue() : null)
                            .seed(seed)
                            .build();
            case "gemini" ->
                    GoogleAiGeminiStreamingChatModel.builder()
                            .apiKey(System.getenv("GEMINI_API_KEY"))
                            .modelName(modelName)
                            .temperature(temperature)
                            .topK(topK)
                            .topP(topP)
                            .logRequestsAndResponses(true)
                            .build();
            default -> throw new IllegalArgumentException("Unsupported provider: " + provider);
        };
    }

    public String baseUrl() {
        return SpringContextHolder.getBean(OllamaProperties.class).getBaseUrl();
    }

}
